package cn.ichiva.dl4j.gnn.common;

import org.nd4j.linalg.api.iter.NdIndexIterator;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;

/**
 * 进化
 */
public interface Evolution {

    /**
     * 遗传
     * @param mDna
     * @param fDna
     * @return
     */
    INDArray inheritance(INDArray mDna,INDArray fDna);

    /**
     * 突变
     * @param dna
     * @param v
     * @param r 突变范围
     * @return
     */
    INDArray mutation(INDArray dna,double v, double r);

    /**
     * 置换
     * @param dna
     * @param v
     * @return
     */
    INDArray substitution(INDArray dna,double v);

    /**
     * 外源
     * @param dna
     * @param v
     * @return
     */
    INDArray other(INDArray dna,double v);

    /**
     * DNA 是否同源
     * @param mDna
     * @param fDna
     * @return
     */
    boolean iSogeny(INDArray mDna, INDArray fDna);

    Evolution defInstance = new DefaultEvolution();

    /**
     * 通用默认实现
     */
    class DefaultEvolution implements Evolution{
        @Override
        public INDArray inheritance(INDArray mDna, INDArray fDna) {
            if(mDna == fDna) return mDna;

            long[] mShape = mDna.shape();
            if(!iSogeny(mDna,fDna)){
                throw new RuntimeException("非同源dna");
            }

            INDArray nDna = Nd4j.create(mShape);
            NdIndexIterator it = new NdIndexIterator(mShape);
            while (it.hasNext()){
                long[] next = it.next();

                double val;
                if(Math.random() > 0.5){
                    val = fDna.getDouble(next);
                }else {
                    val = mDna.getDouble(next);
                }

                nDna.putScalar(next,val);
            }

            return nDna;
        }

        @Override
        public INDArray mutation(INDArray dna, double v, double r) {
            long[] shape = dna.shape();
            INDArray nDna = Nd4j.create(shape);
            NdIndexIterator it = new NdIndexIterator(shape);
            while (it.hasNext()) {
                long[] next = it.next();
                if(Math.random() < v){
                    dna.putScalar(next,dna.getDouble(next) + ((Math.random() - 0.5) * r * 2));
                }else {
                    nDna.putScalar(next,dna.getDouble(next));
                }
            }

            return nDna;
        }

        @Override
        public INDArray substitution(INDArray dna, double v) {
            long[] shape = dna.shape();
            INDArray nDna = Nd4j.create(shape);
            NdIndexIterator it = new NdIndexIterator(shape);
            while (it.hasNext()) {
                long[] next = it.next();
                if(Math.random() > v){
                    long[] tag = new long[shape.length];
                    for (int i = 0; i < shape.length; i++) {
                        tag[i] = (long) (Math.random() * shape[i]);
                    }

                    nDna.putScalar(next,dna.getDouble(tag));
                }else {
                    nDna.putScalar(next,dna.getDouble(next));
                }
            }

            return nDna;
        }

        @Override
        public INDArray other(INDArray dna, double v) {
            long[] shape = dna.shape();
            INDArray nDna = Nd4j.create(shape);
            NdIndexIterator it = new NdIndexIterator(shape);
            while (it.hasNext()) {
                long[] next = it.next();
                if(Math.random() > v){
                    nDna.putScalar(next,Math.random());
                }else {
                    nDna.putScalar(next,dna.getDouble(next));
                }
            }

            return nDna;
        }

        @Override
        public boolean iSogeny(INDArray mDna, INDArray fDna) {
            long[] mShape = mDna.shape();
            long[] fShape = fDna.shape();

            if (mShape.length == fShape.length) {
                for (int i = 0; i < mShape.length; i++) {
                    if (mShape[i] != fShape[i]) {
                        return false;
                    }
                }

                return true;
            }

            return false;
        }
    }
}
